from stack import LinkListStack
from queue import LinkListQueue


class BinarySearchTree:
    class Node:
        def __init__(self, e):
            self.e = e
            self.left = None
            self.right = None

    def __init__(self):
        self._root = None
        self._size = 0

    def size(self):
        return self._size

    def is_empty(self):
        return self._size == 0

    # 插入数据(外部调用)(递归)
    def add(self, e):
        self._root = self._add(self._root, e)

    # 以传入的node为根添加新元素
    # 返回一个新的树以返回的node为根
    def _add(self, node, e):
        # 如果node是空,说明应该插入在这里,生成新node并返回回去让父节点可以连接上
        if node is None:
            self._size += 1
            return self.Node(e)

        if e < node.e:
            # 如果e的值比当前node的值小, 那么从node的左孩子继续递归
            node.left = self._add(node.left, e)
        elif e > node.e:
            # 如果e的值比当前node的值大, 那么从node的右孩子继续递归
            node.right = self._add(node.right, e)
        # 如果e==当前node的e,什么都不做返回
        return node

    # 插入数据(外部调用)(循环)
    def cycle_add(self, e):
        n = self._root
        if self._root is None:
            # 如果根节点为空,直接根据e生成根节点
            self._root = self.Node(e)
            self._size += 1
            return

        while n.e != e:
            if e < n.e:
                if n.left is None:
                    n.left = self.Node(e)
                    self._size += 1
                    break
                else:
                    n = n.left
            elif e > n.e:
                if n.right is None:
                    n.right = self.Node(e)
                    self._size += 1
                    break
                else:
                    n = n.right

    def contains(self, e):
        pass

    # 前序遍历
    def pre_order(self):
        self._pre_order(self._root)

    def _pre_order(self, node):
        if node is None:
            return
        print(node.e, end=',')
        self._pre_order(node.left)
        self._pre_order(node.right)

    # 中序遍历
    def in_order(self):
        self._in_order(self._root)

    def _in_order(self, node):
        if node is None:
            return
        self._in_order(node.left)
        print(node.e, end=',')
        self._in_order(node.right)

    # 后序遍历
    def post_order(self):
        self._post_order(self._root)

    def _post_order(self, node):
        if node is None:
            return
        self._post_order(node.left)
        self._post_order(node.right)
        print(node.e, end=',')

    # 非递归的前序遍历
    def pre_order_not_recursive(self):
        s = LinkListStack()
        s.push(self._root)

        while not s.is_empty():
            node = s.pop()
            print(node.e, end=',')

            # 前序遍历先左再右，但是栈是先入后出，所以先push右节点
            if node.right is not None:
                s.push(node.right)
            if node.left is not None:
                s.push(node.left)

    # 层序遍历
    def level_order(self):
        q = LinkListQueue()
        q.enqueue(self._root)

        while not q.is_empty():
            node = q.dequeue()
            print(node.e, end=',')

            if node.left is not None:
                q.enqueue(node.left)
            if node.right is not None:
                q.enqueue(node.right)

    # 返回树中最小值
    def mini_node(self):
        if self.is_empty():
            raise Exception("Error: 二叉树为空")
        return self._mini_node(self._root)

    def _mini_node(self, node: 'Node'):
        if node.left is None:
            return node
        return self._mini_node(node.left)

    # 返回树中最大值
    def max_node(self):
        if self.is_empty():
            raise Exception("Error: 二叉树为空")
        return self._max_node(self._root)

    def _max_node(self, node: 'Node'):
        if node.right is None:
            return node
        return self._max_node(node.right)

    # 移除树中最小值
    def remove_mini_node(self):
        node = self.mini_node()
        self._root = self._remove_mini_node(self._root)
        return node

    def _remove_mini_node(self, node: 'Node'):
        # 当node的左子树为空时,返回右子树
        if node.left is None:
            self._size -= 1
            right_node = node.right
            # 删除node要释放node.right指针
            node.right = None
            return right_node
        # 右子树绑定到父节点的左子树上,父节点左子树就被移除了
        node.left = self._remove_mini_node(node.left)
        # 将左子树返回给父节点
        return node

    # 移除树中最大值
    def remove_max_node(self):
        node = self.max_node()
        self._root = self._remove_max_node(self._root)
        return node

    def _remove_max_node(self, node: 'Node'):
        # 当node的右子树为空时,返回左子树
        if node.right is None:
            self._size -= 1
            left_node = node.left
            # 删除node要释放node.left指针
            node.left = None
            return left_node
        # 左子树绑定到父节点的右子树上,父节点右子树就被移除了
        node.right = self._remove_max_node(node.right)
        # 将右子树返回给父节点
        return node

    # 移除元素
    def remove_element(self, element):
        if self.is_empty():
            raise Exception("Error: 二叉树为空")
        self._root = self._remove_element(self._root, element)

    def _remove_element(self, node: 'Node', element):
        if node is None:
            return None

        if element < node.e:
            node.left = self._remove_element(node.left, element)
            return node
        if element > node.e:
            node.right = self._remove_element(node.right, element)
            return node

        # element = node.e
        if node.left is None:
            self._size -= 1
            right_node = node.right
            # 要释放node.right指针
            node.right = None
            return right_node
        if node.right is None:
            self._size -= 1
            left_node = node.left
            node.left = None
            return left_node
        # 即有左孩子也有右孩子的情况
        # successor后继,指定为比当前node值大一点的值,位于当前node的右子树最小值
        successor = self._mini_node(node.right)
        successor.right = self._remove_mini_node(node.right)
        # 因为_remove_mini_element中有size-=1就不用在此代码段里写size-=1
        successor.left = node.left

        # 因为删除node, 必须把node的left right指针都指向空
        node.left = None
        node.right = None
        return successor


if __name__ == '__main__':
    l1 = [8, 4, 2, 9, 6, 3, 5, 7, 1]
    #          8
    #      4       9
    #   2     6
    # 1   3  5   7

    b1 = BinarySearchTree()
    for e in l1:
        b1.add(e)

    print("前序遍历")
    b1.pre_order()
    print("\n")

    print("非递归前序遍历")
    b1.pre_order_not_recursive()
    print("\n")

    print("中序遍历")
    b1.in_order()
    print("\n")

    print("后序遍历")
    b1.post_order()
    print("\n")

    print("层序遍历")
    b1.level_order()
    print("\n")

    print(f"b1移除掉3这个node")
    b1.remove_element(4)
    b1.level_order()
    print("\n")


